from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import datetime  # For datetime objects
import os.path  # To manage paths
import sys  # To find out the script name (in argv[0])
sys.path.append('./backtrader')
# data_path_global = './backtrader/datas/orcl-1995-2014.txt'
data_path_global = './csv_data/SH_000300_stock_data.csv'
# Import the backtrader platform
import backtrader as bt


# Create a Stratey
class TestStrategy(bt.Strategy):
    params = (
        ('maperiod_2', 2),
        ('maperiod_5', 5),
        ('maperiod_10', 10),
        ('maperiod_20', 20),
        ('maperiod_30', 30),
        ('maperiod_60', 60),
        ('maperiod_120', 120),
        ('divergence_period', 90),  # 底背离判断的天数
        ('top_divergence_period', 60),  # 顶背离判断的天数
        ('sell_ma_period', 10),  # 卖出信号均线天数
        ('profit_take_pct', 5.0),  # 止盈百分比（3%）
        ('bollinger_period', 60),  # 布林带周期
        ('bollinger_dev', 1.05),    # 布林带标准差倍数
    )

    def log(self, txt, dt=None):
        ''' Logging function fot this strategy'''
        dt = dt or self.datas[0].datetime.date(0)
        print('%s, %s' % (dt.isoformat(), txt))

    def __init__(self):
        # Keep a reference to the "close" line in the data[0] dataseries
        self.dataclose = self.datas[0].close
        # 引用成交量数据
        self.volume = self.datas[0].volume
       
        # To keep track of pending orders and buy price/commission
        self.order = None
        self.buyprice = None
        self.buycomm = None
        
        # 统计变量
        self.trade_count = 0  # 交易次数
        self.start_cash = None  # 初始资金
        self.start_date = None  # 开始日期
        self.end_date = None  # 结束日期

        # Add a MovingAverageSimple indicator
        self.sma = bt.indicators.SimpleMovingAverage(
            self.datas[0], period=self.params.maperiod_5)
        self.sma10 = bt.indicators.SimpleMovingAverage(
            self.datas[0], period=self.params.maperiod_10)
        self.sma20 = bt.indicators.SimpleMovingAverage(
            self.datas[0], period=self.params.maperiod_20)
        self.sma30 = bt.indicators.SimpleMovingAverage(
            self.datas[0], period=self.params.maperiod_30)
        # 添加120日均线
        self.sma120 = bt.indicators.SimpleMovingAverage(
            self.datas[0], period=self.params.maperiod_120)
        # 添加5日线上穿120日线的交叉指标
        self.crossover_5_120 = bt.ind.CrossOver(self.sma, self.sma120)

        
        # Indicators for the plotting show
        # bt.indicators.ExponentialMovingAverage(self.datas[0], period=25)
        # bt.indicators.WeightedMovingAverage(self.datas[0], period=25,
        #                                     subplot=True)
        # bt.indicators.StochasticSlow(self.datas[0])
        self.macd_histo = bt.indicators.MACDHisto(self.datas[0])
        # add diff and dea indicator
        self.macd = bt.indicators.MACD(self.datas[0])
        # MACD包含了三个值：macd (DIFF线), signal (DEA线), histo (柱状图)
        self.diff = self.macd.macd  # DIFF线 (快速EMA - 慢速EMA)
        self.dea = self.macd.signal  # DEA线 (DIFF的EMA平滑线)
        
        # 添加最低价指标用于底背离判断
        self.lowest_period = bt.indicators.Lowest(self.datas[0].low, period=self.params.divergence_period)
        
        # 添加最高价指标用于顶背离判断
        self.highest_period = bt.indicators.Highest(self.datas[0].high, period=self.params.top_divergence_period)
        
        # 添加卖出信号用的均线指标
        self.sell_ma = bt.indicators.SimpleMovingAverage(
            self.datas[0], period=self.params.sell_ma_period)
        
        # self.macd = bt.indicators.MACD(self.datas[0])
        # rsi = bt.indicators.RSI(self.datas[0])
        # bt.indicators.SmoothedMovingAverage(rsi, period=10)
        # bt.indicators.ATR(self.datas[0], plot=False)
        
        # 添加布林带指标
        self.bollinger = bt.indicators.BollingerBands(
            self.datas[0],
            period=self.params.bollinger_period,
            devfactor=self.params.bollinger_dev
        )

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            # Buy/Sell order submitted/accepted to/by broker - Nothing to do
            return

        # Check if an order has been completed
        # Attention: broker could reject order if not enough cash
        if order.status in [order.Completed]:
            if order.isbuy():
                self.log(
                    'BUY EXECUTED, Price: %.2f, Cost: %.2f, Comm %.2f' %
                    (order.executed.price,
                     order.executed.value,
                     order.executed.comm))

                self.buyprice = order.executed.price
                self.buycomm = order.executed.comm
            else:  # Sell
                self.log('SELL EXECUTED, Price: %.2f, Cost: %.2f, Comm %.2f' %
                         (order.executed.price,
                          order.executed.value,
                          order.executed.comm))

            self.bar_executed = len(self)

        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('Order Canceled/Margin/Rejected')

        self.order = None

    def notify_trade(self, trade):
        if not trade.isclosed:
            return

        # 统计交易次数
        self.trade_count += 1
        
        self.log('OPERATION PROFIT, GROSS %.2f, NET %.2f' %
                 (trade.pnl, trade.pnlcomm))


    def start(self):
        # 记录开始信息
        self.start_cash = self.broker.getvalue()
        # 暂时将start_date设为None，稍后在next方法中设置
        self.start_date = None
        self.log('策略开始运行 - 初始资金: %.2f' % self.start_cash)
        
    def stop(self):
        # 记录结束信息
        self.end_date = self.datas[0].datetime.date(0)
        final_cash = self.broker.getvalue()
        self.log('策略结束运行 - 结束日期: %s, 最终资金: %.2f' % (self.end_date, final_cash))
        
        # 计算统计数据
        total_return = ((final_cash - self.start_cash) / self.start_cash) * 100
        
        # 计算交易周期
        if self.start_date and self.end_date:
            trade_period = (self.end_date - self.start_date).days
            # 确保交易周期不为负数
            if trade_period < 0:
                trade_period = 0
        else:
            # 如果没有有效的开始和结束日期，使用默认值0
            trade_period = 0
        
        self.log('='*50)
        self.log('策略统计结果:')
        self.log('交易周期: %s 至 %s (共%d天)' % (self.start_date, self.end_date, trade_period))
        self.log('初始资金: %.2f' % self.start_cash)
        self.log('最终资金: %.2f' % final_cash)
        self.log('总收益率: %.2f%%' % total_return)
        self.log('交易次数: %d' % self.trade_count)
        if trade_period > 0:
            self.log('年化收益率: %.2f%%' % (total_return * 365 / trade_period))
        self.log('='*50)
    def check_buymacd(self):
        # 判断DIFF和DEA在零轴下方形成金叉的条件
        # 金叉：DIFF线从下方穿越DEA线
        # 零轴下方：当前DIFF和DEA都小于0
        # 需要至少2个数据点来判断穿越
        if (len(self) > 1 and 
            self.diff[0] < 0 and self.dea[0] < 0 and  # 当前都在零轴下方
            self.diff[0] > self.dea[0] and  # 当前DIFF在DEA上方
            self.diff[-1] <= self.dea[-1]):  # 前一根K线DIFF在DEA下方或相等
            # self.log('*** 金叉信号：DIFF线在零轴下方穿越DEA线！DIFF: %.2f, DEA: %.2f ***' % 
            #         (self.diff[0], self.dea[0]))
            # self.order = self.buy()
            return  1
        return 0
    def check_sell_macd(self):
        # 判断DIFF和DEA在零轴上方形成死叉的条件
        # 死叉：DIFF线从上方穿越DEA线
        # 零轴上方：当前DIFF和DEA都大于0
        # 需要至少2个数据点来判断穿越
        if (len(self) > 1 and 
            self.diff[0] > 0 and self.dea[0] > 0 and  # 当前都在零轴上方
            self.diff[0] < self.dea[0] and  # 当前DIFF在DEA下方
            self.diff[-1] >= self.dea[-1]):  # 前一根K线DIFF在DEA上方或相等
            # self.log('*** 死叉信号：DIFF线在零轴上方穿越DEA线！DIFF: %.2f, DEA: %.2f ***' % 
            #         (self.diff[0], self.dea[0]))
            # self.order = self.sell()
            return  1
        return 0
    def check_buymacd_diff(self):
        if (len(self) > 1 and 
            self.diff[0] > 0 and  # 当前DIFF在零轴上方
            self.diff[-1] <= 0):  # 前一根K线DIFF在零轴下方或等于零轴
            # self.log('*** 强势信号：DIFF线上穿零轴！DIFF: %.2f, DEA: %.2f ***' % 
            #         (self.diff[0], self.dea[0]))
            # self.order = self.buy()
            return  1
        return 0
    def check_buy_Deviation(self):
        """底背离买入判断函数
        当价格达到指定期间最低价但DIFF没有创新低时（底背离）执行买入
        """
        # 需要至少指定天数的数据点来计算最低价
        if len(self) >= self.params.divergence_period:
            # 当前价格是否等于指定期间最低价
            current_low = self.datas[0].low[0]
            lowest_period_value = self.lowest_period[0]
            
            # 容忍小误差，判断当前低价是否接近指定期间最低价
            if abs(current_low - lowest_period_value) < 0.01:
                # 计算指定期间内DIFF的最低值
                diff_period_min = float('inf')
                for i in range(min(self.params.divergence_period, len(self))):
                    if self.diff[-i] < diff_period_min:
                        diff_period_min = self.diff[-i]
                
                # 判断当前DIFF是否高于指定期间内的最低值（底背离）
                if self.diff[0] > diff_period_min:
                    # self.log('*** 底背离信号：价格创%d日新低但DIFF未创新低！Price: %.2f, DIFF: %.2f, %d日DIFF最低: %.2f ***' % 
                            # (self.params.divergence_period, current_low, self.diff[0], self.params.divergence_period, diff_period_min))
                    # 如果没有仓位，执行买入操作
                    # if not self.position:
                    # self.order = self.buy()
                    return 1
        return 0
    def check_buy_sma_vol(self):
        """判断成交量3日内递增、5日线上穿120日线、价格回到5日线的买入信号函数"""
        # 需要至少3个数据点来判断成交量递增
        # 需要至少120个数据点来确保120日均线有效
        if len(self) >= 120:
            # 条件1：成交量3日内递增
            # 当前成交量 > 前一天成交量 > 前两天成交量
            volume_increasing = (self.volume[0] > self.volume[-1] and 
                                 self.volume[-1] > self.volume[-2])
            
            # 条件2：5日线上穿120日线（最近几天内发生的金叉）
            # 使用交叉指标检测最近3天内是否发生了金叉
            crossover_occurred = False
            for i in range(3):
                if len(self) > i and self.crossover_5_120[-i] > 0:
                    crossover_occurred = True
                    break
            
            # 条件3：价格回到5日线（当前价格接近5日线）
            # 定义一个较小的阈值来判断价格是否接近5日线
            price_near_sma5 = abs(self.dataclose[0] - self.sma[0]) < self.dataclose[0] * 0.005
            
            # 当三个条件都满足时返回1，否则返回0
            if volume_increasing and crossover_occurred and price_near_sma5:
                return 1
        return 0
    def check_sell_stop_win(self):
        """止盈卖出判断函数
        当盈利超过指定百分比时执行卖出
        """
        # 检查当前持仓
        if self.position and self.buyprice is not None:
            # 获取当前价格
            current_price = self.dataclose[0]
            # 计算当前收益率
            profit_pct = ((current_price - self.buyprice) / self.buyprice) * 100
            # 检查是否达到止盈百分比
            if profit_pct >= self.params.profit_take_pct:

                # self.log('*** 止盈卖出信号：收益率达到 %.2f%%，执行卖出 ***' % profit_pct)
                return 1
        return 0
    def check_sell_Deviation(self):
        """顶背离卖出判断函数
        当价格达到指定期间最高价但DIFF没有创新高时（顶背离）返回1，否则返回0
        """
        # 需要至少指定天数的数据点来计算最高价
        if len(self) >= self.params.top_divergence_period:
            # 当前价格是否接近指定期间最高价（考虑小误差）
            current_high = self.datas[0].high[0]
            highest_period_value = self.highest_period[0]
            
            # 容忍小误差，判断当前高价是否接近指定期间最高价
            if abs(current_high - highest_period_value) < 0.01:
                # 计算指定期间内DIFF的最高值
                diff_period_max = float('-inf')
                for i in range(min(self.params.top_divergence_period, len(self))):
                    if self.diff[-i] > diff_period_max:
                        diff_period_max = self.diff[-i]
                
                # 判断当前DIFF是否低于指定期间内的最高值（顶背离）
                if self.diff[0] < diff_period_max:
                    return 1
        return 0
        
    def check_buy_bollinger(self):
        """基于布林带的买入判断函数
        当价格触碰或跌破布林带下轨时返回1，否则返回0
        """
        # 确保有足够的数据计算布林带指标
        if len(self) >= self.params.bollinger_period:
            # 获取当前价格和布林带下轨值
            current_price = self.dataclose[0]
            lower_band = self.bollinger.lines.bot[0]
            
            # 判断价格是否触碰或跌破布林带下轨
            # 考虑一个小的容忍度，避免因价格微小波动而频繁触发信号
            if current_price <= lower_band * 1.005:
                return 1
        return 0
        
    def check_sell_bollinger(self):
        """基于布林带的卖出判断函数
        当价格触碰或突破布林带上轨时返回1，否则返回0
        """
        # 确保有足够的数据计算布林带指标
        if len(self) >= self.params.bollinger_period:
            # 获取当前价格和布林带上轨值
            current_price = self.dataclose[0]
            upper_band = self.bollinger.lines.top[0]
            
            # 判断价格是否触碰或突破布林带上轨
            # 考虑一个小的容忍度，避免因价格微小波动而频繁触发信号
            if current_price >= upper_band * 0.995:
                return 1
        return 0
    def next(self):
        # 记录第一个交易日作为开始日期
        if self.start_date is None:
            self.start_date = self.datas[0].datetime.date(0)
            self.log('策略开始日期记录: %s' % self.start_date)
            
        # Simply log the closing price of the series from the reference
        # self.log('Close_day: %.2f, sma: %.2f, sma10: %.2f, sma20: %.2f, DIFF: %.2f, DEA: %.2f, macd_histo: %.2f' % 
        #         (self.dataclose[0], self.sma[0], self.sma10[0], self.sma20[0], self.diff[0], self.dea[0], self.macd_histo.histo[0]))
                # DIFF和DEA指标已添加并打印 
        
        
        
        # Check if an order is pending ... if yes, we cannot send a 2nd one
        if self.order:
            return
            
            

        #==================================================================
        if self.check_buymacd(): # 调用金叉买入判断函数
            # self.order = self.buy()
            pass
        if self.check_buymacd_diff(): # 调用DIFF线上穿零轴买入判断函数
            # self.order = self.buy()
            pass
        #==================================================================
        if self.check_buy_Deviation():# 调用底背离买入判断函数
            # self.order = self.buy()
            pass
        #==================================================================
        # 调用均线成交量买入判断函数
        if self.check_buy_sma_vol():
            # self.order = self.buy()
            pass
        #==================================================================
        # 调用布林带买入判断函数
        if self.check_buy_bollinger():
            # self.log('*** 布林带买入信号：价格触碰下轨！Price: %.2f, Lower Band: %.2f ***' % 
            #         (self.dataclose[0], self.bollinger.lines.bot[0]))
            self.order = self.buy()
            pass
        
    

        # ====================卖出操作================
        if self.position:
            # 止盈卖出信号
            if self.check_sell_stop_win():
                # self.order = self.sell()
                # self.log('*** 止盈卖出信号，执行卖出 ***')
                pass
            # 死叉卖出信号
            if self.check_sell_macd():
                # self.log('*** 死叉卖出信号：DIFF线在零轴上方穿越DEA线！DIFF: %.2f, DEA: %.2f ***' % 
                #         (self.diff[0], self.dea[0]))
            #     self.order = self.sell()
                pass
            # 顶背离卖出信号
            if self.check_sell_Deviation():
                # self.log('*** 顶背离卖出信号：价格创%d日新高但DIFF未创新高！Price: %.2f, DIFF: %.2f ***' % 
                #         (self.params.top_divergence_period, self.datas[0].high[0], self.diff[0]))
                # self.order = self.sell()
                pass
            # 布林带卖出信号
            if self.check_sell_bollinger():
                # self.log('*** 布林带卖出信号：价格触碰上轨！Price: %.2f, Upper Band: %.2f ***' % 
                #         (self.dataclose[0], self.bollinger.lines.top[0]))
                self.order = self.sell()
                pass
            # # 价格超过指定日均线卖出信号
            # elif self.dataclose[0] > self.sell_ma[0]:
            #     # self.log('*** 卖出信号：价格超过%d日均线！Price: %.2f, MA%d: %.2f ***' % 
            #             # (self.params.sell_ma_period, self.dataclose[0], self.params.sell_ma_period, self.sell_ma[0]))
            #     self.order = self.sell()

        # Check if we are in the market
        # if not self.position:

        #     # Not yet ... we MIGHT BUY if ...
        #     if self.dataclose[0] > self.sma[0]:

        #         # BUY, BUY, BUY!!! (with all possible default parameters)
        #         self.log('BUY CREATE, %.2f' % self.dataclose[0])

        #         # Keep track of the created order to avoid a 2nd order
        #         self.order = self.buy()

        # else:

        #     if self.dataclose[0] < self.sma[0]:
        #         # SELL, SELL, SELL!!! (with all possible default parameters)
        #         self.log('SELL CREATE, %.2f' % self.dataclose[0])

        #         # Keep track of the created order to avoid a 2nd order
        #         self.order = self.sell()


if __name__ == '__main__':
    # Create a cerebro entity
    cerebro = bt.Cerebro()

    # Add a strategy
    cerebro.addstrategy(TestStrategy)

    # Datas are in a subfolder of the samples. Need to find where the script is
    # because it could have been called from anywhere
    modpath = os.path.dirname(os.path.abspath(sys.argv[0]))
    datapath = os.path.join(modpath, data_path_global)

    # Create a Data Feed
    data = bt.feeds.YahooFinanceCSVData(
        dataname=datapath,
        # Do not pass values before this date
        fromdate=datetime.datetime(2023, 8, 3),
        # Do not pass values after this date
        todate=datetime.datetime(2025, 9, 5),
        # Do not pass values after this date
        reverse=False)

    # Add the Data Feed to Cerebro
    cerebro.adddata(data)

    # Set our desired cash start
    cerebro.broker.setcash(10000000.0)

    # Add a FixedSize sizer according to the stake
    cerebro.addsizer(bt.sizers.FixedSize, stake=10)

    # Set the commission
    cerebro.broker.setcommission(commission=0.05)

    # Print out the starting conditions
    print('Starting Portfolio Value: %.2f' % cerebro.broker.getvalue())

    # Run over everything
    cerebro.run()

    # Print out the final result
    final_value = cerebro.broker.getvalue()
    print('Final Portfolio Value: %.2f' % final_value)
    
    
    # Plot the result
    # cerebro.plot()